import torch
from typing import TypeVar
import torch.nn.functional as F

from .impl_van import VAN
from .impl_uper import LikeUPerHead
from .impl_fcn import FCNHead
from .impl_cls import ClassifyHeader
from .config import VAN_ARGS, FCN_ARGS, UPER_ARGS
from ..interface import Module

T = TypeVar('T', bound=torch.nn.Module)


class Net(Module):
    model_id = 'jassor'

    def __init__(self, pretrained: str):
        super().__init__()
        # 使用 VAN 原定的编码器以 1 的维度输入，并且初始化权重
        self.encoder = VAN(**{**VAN_ARGS, 'in_chans': 3})
        self.encoder.init_weights(pretrained=pretrained)

        # 使用 VAN 原定的解码器做分割网络结果
        self.decoder = LikeUPerHead(**{**UPER_ARGS, 'out_dim': 5})
        # 以及使用 VAN 原定的辅助训练网络帮助分割任务的训练
        self.auxiliary = FCNHead(**{**FCN_ARGS, 'num_classes': 5})
        self.activation = torch.nn.Softmax(dim=1)

        # 上面全是分割网络的基本结构，下面开始就是自己作妖设计的伴随分类网络了
        self._mode = 'train_seg'
        self.classify = ClassifyHeader()
        self.device = 'cpu'

    def to(self, device):
        self.device = device
        return super().to(device)

    def mode(self, mode: str = None) -> str:
        if mode and mode == 'train_seg':
            # 训练分割的时候，解锁分割网络梯度
            self._mode = mode
            self.train()
        elif mode and mode == 'train_cls':
            # 训练分类的时候，锁定分割网络梯度
            self._mode = mode
            self.encoder.eval()
            self.decoder.eval()
            self.classify.train()
        elif mode and mode == 'predict':
            self._mode = mode
            self.eval()
        return self._mode

    def forward(self, image: torch.tensor):
        # -------------------- 分割网络部分 ---------------------------
        h, w = image.shape[2:]
        # 当训练分割网络时，计算和返回辅助网络结果
        if self._mode == 'train_seg':
            image = image.to(self.device)
            # features 会被用来分类
            features = self.encoder(image)
            # seg_result 是真实的网络输出
            seg_result = self.decoder(features)
            seg_result = F.interpolate(seg_result, size=(h, w), mode='bilinear', align_corners=False)
            seg_result = self.activation(seg_result)
            # aux_result 是辅助训练网络
            aux_result = self.auxiliary(features)
            aux_result = F.interpolate(aux_result, size=(h, w), mode='bilinear', align_corners=False)
            aux_result = self.activation(aux_result)
            return [seg_result, aux_result]
        # -------------------- 分类网络部分 ---------------------------

        # features 会被用来分类
        all_features = []
        seg_results = []
        # !!!!!!!!!!!!!!!!!!!!!!!!
        batch = 1
        for b in range(0, image.shape[0], batch):
            features = self.encoder(image[b: b + batch, :, :, :].to(self.device))
            features = [feature.detach() for feature in features]
            seg_result = self.activation(self.decoder(features))
            all_features.append(torch.concat([
                torch.max_pool2d(
                    feature * torch.max_pool2d(seg_result[:, seg_channel: seg_channel + 1, :, :], kernel_size=2 ** index),
                    kernel_size=2 ** (3 - index)
                ).detach()
                for seg_channel in range(seg_result.shape[1]) for index, feature in enumerate(features)
            ], dim=1))
            seg_result = seg_result.detach().cpu()
            seg_results.append(seg_result)
            del features, seg_result
        features = torch.concat(all_features, dim=0).detach()
        # 然后将 batch 维度转换过去变成 3D 结构，batch 维度强制缩为 1，channel、h、w 保持逻辑不变
        features = features[None, ...].permute(0, 2, 1, 3, 4)
        seg_result = torch.concat(seg_results, dim=0)
        del all_features, seg_results

        if self._mode == 'train_cls':
            cls_result = self.classify(features)
            return cls_result
        # -------------------- 预测部分 ---------------------------
        if self._mode == 'predict':
            cls_result = self.classify(features).detach().cpu()
            seg_result = F.interpolate(seg_result, size=(h, w), mode='bilinear', align_corners=False).detach().cpu()
            return [seg_result, cls_result]
        raise Exception(f'mode not available as "{self._mode}" not in [train_seg, train_cls, predict]')
